package main

import (
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"os/exec"
	"strconv"
	"sync"
	"time"

	"github.com/gorilla/websocket"
)

var upgrader = websocket.Upgrader{
	ReadBufferSize:  8192,
	WriteBufferSize: 8192,
	CheckOrigin:     func(r *http.Request) bool { return true },
}

// 流管理器
type StreamManager struct {
	streams   map[string]*Stream // RTSP URL到Stream的映射
	streamIDs map[string]string  // streamID到RTSP URL的映射
	mutex     sync.RWMutex
}

// 客户端结构
type Client struct {
	conn          *websocket.Conn
	sendQueue     chan []byte
	closed        bool
	mutex         sync.Mutex
	streamID      string
	rtspURL       string
	lastActive    time.Time
	queueCapacity int  // 队列容量
	sentStartup   bool // 是否已发送 startupCache
}

// 流结构
type Stream struct {
	rtspURL      string
	cmd          *exec.Cmd
	stdout       io.ReadCloser
	clients      map[*Client]bool // 客户端列表
	clientCount  int              // 客户端计数
	maxClients   int              // 最大客户端数
	mutex        sync.RWMutex     // 客户端操作锁
	stopped      bool             // 停止标志
	stopOnce     sync.Once        // 确保stop只执行一次
	dataQueue    chan []byte      // 视频数据队列
	closeChan    chan struct{}    // 关闭通知
	healthCheck  chan error       // 健康检查通道
	startupCache [][]byte         // FLV 起始数据缓存
}

// 创建新的流管理器
func NewStreamManager() *StreamManager {
	return &StreamManager{
		streams:   make(map[string]*Stream),
		streamIDs: make(map[string]string),
	}
}

// 处理WebSocket连接
func (sm *StreamManager) handleWebSocket(w http.ResponseWriter, r *http.Request) {
	// 从URL中获取streamID和rtspURL参数
	streamID := r.URL.Query().Get("streamID")
	if streamID == "" {
		http.Error(w, "Missing streamID parameter", http.StatusBadRequest)
		return
	}

	rtspURL := r.URL.Query().Get("rtspURL")
	if rtspURL == "" {
		// 大华H.265流的默认RTSP URL模板
		//下面是使用ffmpeg保存mp4的方法
		//ffmpeg -rtsp_transport tcp   -i "rtsp://admin:admin123456@101.201.62.217:554/cam/realmonitor?channel=1&subtype=1"  -an -c:v copy output.mp4
		//rtsp://admin:admin123456@101.201.62.217:554/cam/realmonitor?channel=1&subtype=0
		//rtsp://admin:admin123456@101.201.62.217:554/cam/realmonitor?channel=1&subtype=1
		rtspURL = fmt.Sprintf("rtsp://admin:admin123456@192.168.1.108:554/cam/realmonitor?channel=1&subtype=1")
		//rtspURL = fmt.Sprintf("rtsp://admin:admin123456@101.201.62.217:554/cam/realmonitor?channel=1&subtype=1")
	}

	// 获取最大客户端数参数，默认为50
	maxClientsStr := r.URL.Query().Get("maxClients")
	maxClients := 50
	if maxClientsStr != "" {
		if val, err := strconv.Atoi(maxClientsStr); err == nil {
			maxClients = val
		}
	}

	log.Printf("New connection - streamID: %s, normalized RTSP URL: %s, maxClients: %d",
		streamID, rtspURL, maxClients)

	// 升级WebSocket
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Printf("WebSocket upgrade error: %v", err)
		return
	}

	// 创建客户端
	client := &Client{
		conn:          conn,
		sendQueue:     make(chan []byte, 2000),
		closed:        false,
		streamID:      streamID,
		rtspURL:       rtspURL,
		lastActive:    time.Now(),
		queueCapacity: 2000,
	}

	// 设置连接处理器
	conn.SetCloseHandler(func(code int, text string) error {
		client.close()
		sm.mutex.Lock()
		defer sm.mutex.Unlock()
		log.Printf("WebSocket closed from client %s", streamID)
		// 查找对应的RTSP URL
		rtspURL, exists := sm.streamIDs[streamID]
		if !exists {
			return nil
		}

		// 从streamID映射中删除
		delete(sm.streamIDs, streamID)

		// 查找对应的流
		if stream, exists := sm.streams[rtspURL]; exists {
			stream.removeClient(client)
		}
		return nil
	})

	// 管理流创建
	sm.mutex.Lock()
	defer sm.mutex.Unlock()

	// 检查是否已有该RTSP URL的流
	var streamExists bool
	var stream *Stream
	if existing, exists := sm.streams[rtspURL]; exists {
		if existing.stopped {
			log.Printf("移除已经停止的流: %s", rtspURL)
			delete(sm.streams, rtspURL)
		} else {
			stream = existing
			streamExists = true
			log.Printf("重用现有流 - RTSP URL: %s, clients: %d",
				rtspURL, stream.clientCount)

			// 检查是否超过最大客户端数
			if stream.clientCount >= stream.maxClients {
				log.Printf("Max client limit reached for stream: %s, clients: %d",
					rtspURL, stream.clientCount)
				conn.Close()
				return
			}
		}
	}

	// 流不存在则创建新流
	if !streamExists {
		stream = &Stream{
			rtspURL:     rtspURL,
			clients:     make(map[*Client]bool),
			clientCount: 0,
			maxClients:  maxClients,
			stopped:     false,
			stopOnce:    sync.Once{},
			dataQueue:   make(chan []byte, 100), // 视频数据队列
			closeChan:   make(chan struct{}),    // 关闭通知
			healthCheck: make(chan error, 1),    // 健康检查通道
		}

		// 启动FFmpeg，优化参数
		log.Printf("Creating new stream - RTSP URL: %s", rtspURL)
		if err := stream.startFFmpeg(rtspURL); err != nil {
			log.Printf("Failed to start FFmpeg: %v", err)
			conn.Close()
			return
		}

		// 启动数据读取、分发和健康检查协程
		go stream.readFFmpegData()
		go stream.distributeData()
		go stream.healthCheckFFmpeg()

		sm.streams[rtspURL] = stream
		log.Printf("新的rtsp流转换创建 - RTSP URL: %s, maxClients: %d", rtspURL, maxClients)
	} else if stream.clientCount >= stream.maxClients {
		log.Printf("已达到最大客户端数量限制: %s, clients: %d",
			rtspURL, stream.clientCount)
		conn.Close()
		return
	}

	// 记录streamID到RTSP URL的映射
	sm.streamIDs[streamID] = rtspURL

	// 添加客户端到流
	stream.mutex.Lock()
	stream.clients[client] = true
	stream.clientCount++
	log.Printf("Client added - streamID: %s, RTSP URL: %s, total clients: %d",
		streamID, rtspURL, stream.clientCount)
	stream.mutex.Unlock()

	// 启动写入和保活协程
	go client.writeLoop()
	go client.keepAlive()
}

// 客户端写入循环 Client.writeLoop()
func (c *Client) writeLoop() {
	defer func() {
		c.close()
	}()
	for data := range c.sendQueue {
		if len(data) == 1 && data[0] == websocket.PingMessage {
			c.conn.WriteMessage(websocket.PingMessage, nil)
		} else if len(data) == 1 && data[0] == websocket.PongMessage {
			c.conn.WriteControl(websocket.PongMessage, []byte{}, time.Now().Add(time.Second))
		} else {
			err := c.conn.WriteMessage(websocket.BinaryMessage, data)
			if err != nil {
				log.Printf("WebSocket write error: %v", err)
				break // 退出循环会触发 defer 调用 close
			}
		}
		// 更新最后活跃时间
		c.lastActive = time.Now()
	}
}

// 保持连接活跃
func (c *Client) keepAlive() {
	ticker := time.NewTicker(10 * time.Second)
	defer ticker.Stop()
	for {
		select {
		case <-ticker.C:
			c.mutex.Lock()
			if c.closed {
				c.mutex.Unlock()
				return
			}
			c.mutex.Unlock()
			select {
			case c.sendQueue <- []byte{websocket.PingMessage}:
				c.lastActive = time.Now() // 更新活跃时间
			default:
				log.Printf("Ping queue full, disconnecting client - streamID: %s", c.streamID)
				c.close()
				return
			}
		}
	}
}

// 关闭客户端
func (c *Client) close() {
	c.mutex.Lock()
	defer c.mutex.Unlock()
	if !c.closed {
		c.closed = true
		log.Printf("Closing client - streamID: %s, RTSP URL: %s", c.streamID, c.rtspURL)
		c.conn.Close() // 增加这行用于强制关闭底层连接
	}
}

func (c *Client) Send(data []byte) bool {
	c.mutex.Lock()
	defer c.mutex.Unlock()
	if c.closed {
		return false
	}
	// 发送 startupCache
	if !c.sentStartup {
		stream := getStreamByURL(c.rtspURL)
		if stream != nil {
			for _, chunk := range stream.startupCache {
				select {
				case c.sendQueue <- chunk:
				default:
					log.Printf("Startup cache send queue full for client %s", c.streamID)
					return false
				}
			}
		}
		c.sentStartup = true
	}

	select {
	case c.sendQueue <- data:
		return true
	default:
		return false
	}
}

var manager *StreamManager

func getStreamByURL(rtspURL string) *Stream {
	manager := manager
	manager.mutex.RLock()
	defer manager.mutex.RUnlock()
	if stream, ok := manager.streams[rtspURL]; ok {
		return stream
	}
	return nil
}

// 添加 stream.removeClient 方法
func (s *Stream) removeClient(c *Client) {
	s.mutex.Lock()
	defer s.mutex.Unlock()
	if _, ok := s.clients[c]; ok {
		delete(s.clients, c)
		s.clientCount--
		log.Printf("Client removed - streamID: %s, total clients: %d", c.streamID, s.clientCount)
		if s.clientCount <= 0 {
			log.Printf("Stopping stream due to no clients - RTSP URL: %s", s.rtspURL)
			go s.stop()
		}
	}
}

// 启动FFmpeg转码
func (s *Stream) startFFmpeg(rtspURL string) error {
	log.Printf("Starting FFmpeg for stream - RTSP URL: %s", rtspURL)

	cmd := exec.Command(
		"ffmpeg",
		"-hide_banner",
		"-loglevel", "warning",
		"-rtsp_transport", "tcp",
		"-i", rtspURL,
		// 优化参数
		"-c:v", "libx264", // AMD AMF硬件加速
		"-preset", "veryfast", // 平衡质量和速度
		"-b:v", "1500k", // 降低目标码率
		"-maxrate", "2000k", // 降低最大码率
		"-bufsize", "3000k", // 调整缓冲区大小
		"-framerate", "25", // 降低帧率
		"-g", "30", // I帧间隔
		"-keyint_min", "30", // 最小关键帧间隔
		"-sc_threshold", "0", // 禁用场景切换
		"-c:a", "aac", // 将 G.711A 转为 AAC
		"-ar", "44100", // 重采样为 44100Hz（Web兼容）
		"-b:a", "128k", // 设置音频码率
		"-f", "flv", // 输出FLV格式
		"-", // 标准输出
	)

	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return fmt.Errorf("Failed to get FFmpeg stdout: %v", err)
	}

	if err := cmd.Start(); err != nil {
		stdout.Close()
		return fmt.Errorf("Failed to start FFmpeg: %v", err)
	}

	s.cmd = cmd
	s.stdout = stdout
	s.stopped = false
	return nil
}

// 从FFmpeg读取数据
func (s *Stream) readFFmpegData() {
	defer close(s.dataQueue)
	defer s.stop()

	buf := make([]byte, 65536)
	log.Printf("Starting data reader for stream - RTSP URL: %s", s.rtspURL)

	for {
		s.mutex.RLock()
		if s.stopped {
			s.mutex.RUnlock()
			log.Printf("Data reader stopped - RTSP URL: %s", s.rtspURL)
			break
		}
		s.mutex.RUnlock()

		n, err := s.stdout.Read(buf)
		if err != nil {
			if err != io.EOF {
				log.Printf("FFmpeg read error - RTSP URL: %s: %v", s.rtspURL, err)
				s.healthCheck <- err // 通知健康检查
			}
			break
		}

		if n > 0 {
			data := make([]byte, n)
			copy(data, buf[:n])
			// 缓存前面几段作为 startupCache（只缓存前 ~2 秒，估计 200KB）
			if len(s.startupCache) < 30 {
				s.startupCache = append(s.startupCache, data)
			}
			// 将数据放入队列
			select {
			case s.dataQueue <- data:
			case <-s.closeChan:
				log.Printf("Data reader closing - RTSP URL: %s", s.rtspURL)
				return
			}
		}
	}
}

// 分发数据给所有客户端
func (s *Stream) distributeData() {
	log.Printf("Starting data distributor for stream - RTSP URL: %s", s.rtspURL)

	for data := range s.dataQueue {
		// 使用读锁，允许多个goroutine并发读取客户端列表
		s.mutex.RLock()
		clientCount := len(s.clients)
		clientsCopy := make(map[*Client]bool, clientCount)

		for client := range s.clients {
			clientsCopy[client] = true
		}
		s.mutex.RUnlock()

		if clientCount > 0 {
			// 并行处理客户端发送
			var wg sync.WaitGroup
			wg.Add(len(clientsCopy))

			for client := range clientsCopy {
				go func(c *Client) {
					defer wg.Done()

					if time.Since(c.lastActive) > 30*time.Second {
						log.Printf("Client inactive for too long, disconnecting - streamID: %s", c.streamID)
						c.close()
						s.removeClient(c)
						return
					}

					if !c.Send(data) {
						log.Printf("Client queue full or closed, disconnecting - streamID: %s", c.streamID)
						c.close()
						s.removeClient(c)
					}
				}(client)
			}

			wg.Wait() // 等待所有客户端处理完成
		}
	}

	log.Printf("Data distributor stopped - RTSP URL: %s", s.rtspURL)
}

// 健康检查FFmpeg进程
func (s *Stream) healthCheckFFmpeg() {
	log.Printf("Starting health check for stream - RTSP URL: %s", s.rtspURL)

	if err := s.cmd.Wait(); err != nil {
		log.Printf("FFmpeg process terminated - RTSP URL: %s: %v", s.rtspURL, err)

		// 发送健康检查错误
		s.healthCheck <- err
		// 发送健康检查错误（如果需要）
		select {
		case s.healthCheck <- err:
		default:
		}

		// 不直接关闭closeChan，调用stop
		s.stop()
	} else {
		// 正常退出也调用stop
		s.stop()
	}
}

// 停止流
func (s *Stream) stop() {
	s.stopOnce.Do(func() {
		log.Printf("Stopping stream - RTSP URL: %s", s.rtspURL)
		s.mutex.Lock()
		defer s.mutex.Unlock()
		if s.stopped {
			return
		}
		s.stopped = true

		// 关闭通知通道
		close(s.closeChan)

		// 终止FFmpeg进程
		if s.cmd != nil {
			log.Printf("Killing FFmpeg process - RTSP URL: %s", s.rtspURL)
			err := s.cmd.Process.Kill()
			if err != nil {
				log.Printf("FFmpeg kill error: %v", err)
			}
			if err := s.cmd.Wait(); err != nil {
				log.Printf("FFmpeg wait error: %v", err)
			} else {
				log.Printf("FFmpeg exited cleanly.")
			}
			s.cmd = nil
		}

		// 关闭 stdout
		if s.stdout != nil {
			s.stdout.Close()
			s.stdout = nil
		}

		// 清理所有客户端
		for client := range s.clients {
			client.close()
			delete(s.clients, client)
		}
		s.clientCount = 0
	})
}

func (c *Client) IsClosed() bool {
	c.mutex.Lock()
	defer c.mutex.Unlock()
	return c.closed
}

// 系统状态监控接口
func (sm *StreamManager) getSystemStats(w http.ResponseWriter, r *http.Request) {
	sm.mutex.RLock()
	defer sm.mutex.RUnlock()

	w.Header().Set("Content-Type", "application/json")
	response := map[string]interface{}{
		"timestamp":    time.Now().UTC().Format(time.RFC3339),
		"streams":      len(sm.streams),
		"totalClients": 0,
	}

	streams := make(map[string]interface{})
	for rtspURL, stream := range sm.streams {
		stream.mutex.RLock()
		streams[rtspURL] = map[string]interface{}{
			"clientCount": stream.clientCount,
			"maxClients":  stream.maxClients,
			"active":      !stream.stopped,
		}
		stream.mutex.RUnlock()
		response["totalClients"] = response["totalClients"].(int) + stream.clientCount
	}

	response["streamsDetails"] = streams
	if err := json.NewEncoder(w).Encode(response); err != nil {
		http.Error(w, "Failed to encode response", http.StatusInternalServerError)
	}
}

func main() {
	manager = NewStreamManager()
	// 注册WebSocket处理函数
	//ws://127.0.0.1:8082/ws?streamID=123
	http.HandleFunc("/ws", manager.handleWebSocket)
	// 添加系统状态监控接口
	http.HandleFunc("/stats", func(w http.ResponseWriter, r *http.Request) {
		manager.getSystemStats(w, r)
	})
	log.Println("Server started on :8082")
	log.Fatal(http.ListenAndServe(":8082", nil))
}
